package net.cyue.ort.llm.util;

import net.cyue.ort.llm.template.ModelChatMessage;
import net.cyue.ort.llm.template.ModelChatTemplate;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;

/**
 * 提示预处理模块
 * 负责将聊天消息列表转换为模型可接受的提示文本
 */
public class PromptPreprocessor {
    
    private static final Logger logger = LoggerFactory.getLogger(PromptPreprocessor.class);
    
    private static final String DEFAULT_SYSTEM_PROMPT = "SYSTEM:你是一个慷慨解囊的助手。\n\n";
    
    private final ModelChatTemplate chatTemplate;
    private final boolean useDefaultTemplate;
    
    /**
     * 使用指定的chat template创建预处理器
     */
    public PromptPreprocessor(ModelChatTemplate chatTemplate) {
        this.chatTemplate = chatTemplate;
        this.useDefaultTemplate = (chatTemplate == null);
        if (useDefaultTemplate) {
            logger.warn("chatTemplate为空，将使用默认模板格式");
        }
    }
    
    /**
     * 使用默认模板格式创建预处理器
     */
    public static PromptPreprocessor withDefaultTemplate() {
        return new PromptPreprocessor(null);
    }
    
    /**
     * 预处理聊天消息列表，返回格式化后的提示文本
     * 
     * @param messages 聊天消息列表
     * @return 格式化后的提示文本
     */
    public String preprocess(List<ModelChatMessage> messages) {
        if (messages == null || messages.isEmpty()) {
            throw new IllegalArgumentException("消息列表不能为空");
        }
        
        if (useDefaultTemplate) {
            return applyDefaultTemplate(messages);
        }
        
        return chatTemplate.apply(messages);
    }
    
    /**
     * 应用默认模板格式
     */
    private String applyDefaultTemplate(List<ModelChatMessage> messages) {
        StringBuilder sb = new StringBuilder();
        sb.append(DEFAULT_SYSTEM_PROMPT);
        
        for (ModelChatMessage message : messages) {
            sb.append(message.role);
            sb.append(":");
            sb.append(message.content);
            sb.append("\n");
        }
        
        return sb.toString();
    }
    
    /**
     * 检查是否使用了自定义模板
     */
    public boolean hasCustomTemplate() {
        return !useDefaultTemplate;
    }
    
    /**
     * 获取当前使用的chat template（可能为null）
     */
    public ModelChatTemplate getChatTemplate() {
        return chatTemplate;
    }
}

